import json
import logging
import sys
import os
from typing import Any, Dict, List
import requests
import configparser
from datetime import datetime
from tools.utils import BioInfoError, ConfigManager, DataFormatter

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class OpenTargetsManager:
    """OpenTargets API 管理器"""
    
    def __init__(self, config_manager: ConfigManager):
        """初始化OpenTargets管理器"""
        self.config = config_manager
        self.session = requests.Session()
        self.session.headers.update({
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        })
    
    def _make_request(self, method: str, url: str, **kwargs) -> requests.Response:
        """发送HTTP请求并处理错误"""
        try:
            response = self.session.request(method, url, **kwargs)
            response.raise_for_status()
            return response
        except requests.exceptions.RequestException as e:
            raise BioInfoError(f"API请求失败: {e}")
    
    def get_ensembl_id(self, gene_name: str) -> str:
        """根据基因名称获取Ensembl ID"""
        url = f"{self.config.get('API', 'ensembl_api_url')}{gene_name}"
        params = {'content-type': 'application/json'}
        
        response = self._make_request('GET', url, params=params)
        data = response.json()
        
        target_id = data.get("id")
        if not target_id:
            raise BioInfoError(f"未找到基因 '{gene_name}' 的Ensembl ID")
        
        return target_id
    
    def get_disease_efo_id(self, disease_name: str, interactive: bool = True, 
                          for_gradio: bool = False) -> str:
        """根据疾病名称获取EFO ID
        
        Args:
            disease_name: 疾病名称
            interactive: 是否使用交互式选择（当有多个匹配结果时）
            for_gradio: 是否在Gradio界面中使用
            
        Returns:
            疾病的EFO ID 或 (当for_gradio=True且有多个匹配时) 返回("multiple_matches", 匹配结果列表)
        """
        url = self.config.get('API', 'ols_url')
        params = {
            "q": disease_name,
            "ontology": "efo",
            "exact": "true"
        }
        
        response = self._make_request('GET', url, params=params)
        data = response.json()
        
        if data["response"]["numFound"] == 0:
            raise BioInfoError(f"未找到疾病 '{disease_name}' 的EFO ID")
        
        # 过滤并排序疾病结果
        docs = data["response"]["docs"]
        sorted_disease_hits = sorted(
            docs, 
            key=lambda x: x["label"].lower().startswith(disease_name.lower()), 
            reverse=True
        )
        
        if len(sorted_disease_hits) == 1:
            # 只有一个匹配结果，直接返回
            return sorted_disease_hits[0]["iri"].split("/")[-1]
        else:
            # 多个匹配结果
            if for_gradio:
                # 在Gradio界面中，返回特殊标记和匹配列表
                return "multiple_matches", sorted_disease_hits
            elif interactive:
                # 命令行交互式选择
                return self._select_disease_interactive(data, disease_name)
            else:
                # 非交互式模式，返回第一个结果
                logger.warning(f"找到多个匹配的疾病: {disease_name}，使用第一个匹配")
                return sorted_disease_hits[0]["iri"].split("/")[-1]
    
    def _select_disease_interactive(self, data: Dict[str, Any], 
                                   disease_name: str) -> str:
        """交互式选择疾病"""
        docs = data["response"]["docs"]
        print(f"\n找到 {len(docs)} 个匹配的疾病，请选择要查询的疾病：")
        print("-" * 60)
        
        for index, doc in enumerate(docs, start=1):
            efo_id = doc["iri"].split("/")[-1]
            label = doc.get('label', '未知')
            description = doc.get('description', '无描述')
            print(f"{index}. {label}")
            print(f"   EFO ID: {efo_id}")
            if description and description != '无描述':
                print(f"   描述: {description[:100]}{'...' if len(description) > 100 else ''}")
            print()
        
        while True:
            try:
                choice = input(
                    f"请输入要选择的序号 (1-{len(docs)})，或按 q 退出: "
                ).strip()
                if choice.lower() == 'q':
                    raise BioInfoError("用户取消操作")
                
                choice_num = int(choice)
                if 1 <= choice_num <= len(docs):
                    selected_doc = docs[choice_num - 1]
                    selected_id = selected_doc["iri"].split("/")[-1]
                    selected_label = selected_doc.get('label', '未知')
                    print(f"已选择: {selected_label} ({selected_id})")
                    return selected_id
                else:
                    print(f"请输入 1-{len(docs)} 之间的数字")
            except ValueError:
                print("请输入有效的数字或 q 退出")
    
    def send_graphql_request(self, query: str, variables: Dict[str, Any]) -> Dict[str, Any]:
        """发送GraphQL请求"""
        url = self.config.get('API', 'base_url')
        payload = {"query": query, "variables": variables}
        
        response = self._make_request('POST', url, json=payload)
        return response.json()
    
    def query_disease_drugs(self, disease_id: str, limit: int = 100) -> Dict[str, Any]:
        """查询特定疾病的药物信息"""
        # 使用最简单的查询格式
        query = """
        query DiseaseDrugs($diseaseId: String!, $size: Int!) {
          disease(efoId: $diseaseId) {
            id
            name
            knownDrugs(size: $size) {
              count
              rows {
                drug {
                  id
                  name
                }
                phase
                status
              }
            }
          }
        }
        """
        
        variables = {
            "diseaseId": disease_id,
            "size": limit
        }
        return self.send_graphql_request(query, variables)
    
    def query_target_associations(self, gene_id: str, limit: int = 100) -> Dict[str, Any]:
        """查询特定基因关联的疾病"""
        query = """
        query TargetAssociations($id: String!, $size: Int!) {
          target(ensemblId: $id) {
            id
            approvedSymbol
            associatedDiseases(
              page: { index: 0, size: $size }
              orderByScore: "score"
              enableIndirect: false
            ) {
              count
              rows {
                disease { id name }
                score
                datasourceScores { componentId: id score }
              }
            }
          }
        }
        """
        
        variables = {
            "id": gene_id,
            "size": limit
        }
        return self.send_graphql_request(query, variables)
    
    def query_disease_associations(self, disease_id: str, limit: int = 100) -> Dict[str, Any]:
        """查询特定疾病关联的靶点"""
        query = """
        query DiseaseAssociations($id: String!, $size: Int!) {
          disease(efoId: $id) {
            id
            name
            associatedTargets(
              page: { index: 0, size: $size }
              orderByScore: "score"
              enableIndirect: true
            ) {
              count
              rows {
                target {
                  id
                  approvedSymbol
                  approvedName
                }
                score
                datasourceScores { componentId: id score }
              }
            }
          }
        }
        """
        
        variables = {
            "id": disease_id,
            "size": limit
        }
        return self.send_graphql_request(query, variables)


class OpenTargetsDataFormatter:
    """OpenTargets特定的数据格式化工具"""
    
    @staticmethod
    def convert_to_csv(data: Dict[str, Any], filename: str, query_type: str) -> None:
        """将数据转换为CSV格式"""
        csv_filename = filename.replace('.json', '.csv')
        
        try:
            # 确保我们处理的是正确的数据结构
            # 健壮性检查：确保数据是字典类型
            if not isinstance(data, dict):
                raise ValueError(f"期望字典类型的数据，得到: {type(data)}")
            
            with open(csv_filename, 'w', newline='', encoding='utf-8') as csvfile:
                if query_type == 'disease_drugs':
                    OpenTargetsDataFormatter._write_disease_drugs_csv(data, csvfile)
                elif query_type == 'target_associations':
                    OpenTargetsDataFormatter._write_target_associations_csv(data, csvfile)
                elif query_type == 'disease_associations':
                    OpenTargetsDataFormatter._write_disease_associations_csv(data, csvfile)
            
            logger.info(f"CSV结果已保存到: {csv_filename}")
        except Exception as e:
            logger.error(f"CSV转换失败，数据类型: {type(data)}")
            raise BioInfoError(f"保存CSV文件失败: {e}")

    @staticmethod
    def _write_disease_drugs_csv(data: Dict[str, Any], csvfile) -> None:
        """写入疾病药物信息的CSV格式"""
        import csv
        
        try:
            headers = ["drug.id", "drug.name", "phase", "status"]
            writer = csv.writer(csvfile)
            writer.writerow(headers)
            
            # 健壮性处理：确保数据结构正确
            drugs = []
            if isinstance(data, dict) and 'data' in data and isinstance(data['data'], dict):
                disease = data['data'].get('disease', {})
                if isinstance(disease, dict):
                    known_drugs = disease.get('knownDrugs', {})
                    if isinstance(known_drugs, dict):
                        drugs = known_drugs.get('rows', [])
            
            # 确保drugs是列表
            if not isinstance(drugs, list):
                drugs = []
            
            for drug_info in drugs:
                # 确保drug_info是字典
                if not isinstance(drug_info, dict):
                    continue
                    
                drug = drug_info.get('drug', {})
                # 确保drug是字典
                if not isinstance(drug, dict):
                    drug = {}
                    
                writer.writerow([
                    drug.get('id', ''),
                    drug.get('name', ''),
                    drug_info.get('phase', ''),
                    drug_info.get('status', '')
                ])
        except Exception as e:
            logger.error(f"写入疾病药物CSV失败: {str(e)}")
            # 写入错误信息作为CSV的第一行
            writer = csv.writer(csvfile)
            writer.writerow([f"错误: {str(e)}"])
    
    @staticmethod
    def _write_target_associations_csv(data: Dict[str, Any], csvfile) -> None:
        """写入目标关联疾病CSV数据"""
        import csv
        
        try:
            # 获取所有唯一的datasource.id
            rows = []
            if isinstance(data, dict) and 'data' in data and isinstance(data['data'], dict):
                target = data['data'].get('target', {})
                if isinstance(target, dict):
                    associated_diseases = target.get('associatedDiseases', {})
                    if isinstance(associated_diseases, dict):
                        rows = associated_diseases.get('rows', [])
            
            # 确保rows是列表
            if not isinstance(rows, list):
                rows = []
                
            all_datasources = set()
            for row in rows:
                if isinstance(row, dict):
                    scores = row.get('datasourceScores', [])
                    if isinstance(scores, list):
                        for score in scores:
                            if isinstance(score, dict):
                                all_datasources.add(score.get('componentId', ''))
            
            # 排序数据源ID以确保一致的列顺序
            datasource_columns = sorted(all_datasources)
            
            # 构建表头
            headers = ["disease.id", "disease.name", "score"] + datasource_columns
            writer = csv.writer(csvfile)
            writer.writerow(headers)
            
            # 写入数据
            for row in rows:
                if not isinstance(row, dict):
                    continue
                    
                disease = row.get('disease', {})
                if not isinstance(disease, dict):
                    disease = {}
                    
                scores = row.get('datasourceScores', [])
                if not isinstance(scores, list):
                    scores = []
                    
                # 创建数据源分数映射
                score_map = {}
                for score in scores:
                    if isinstance(score, dict):
                        score_map[score.get('componentId', '')] = str(score.get('score', ''))
                
                # 构建行数据
                row_data = [
                    disease.get('id', ''),
                    disease.get('name', ''),
                    str(row.get('score', ''))
                ] + [score_map.get(ds, '') for ds in datasource_columns]
                
                writer.writerow(row_data)
        except Exception as e:
            logger.error(f"写入目标关联疾病CSV失败: {str(e)}")
            # 写入错误信息作为CSV的第一行
            writer = csv.writer(csvfile)
            writer.writerow([f"错误: {str(e)}"])
    
    @staticmethod
    def _write_disease_associations_csv(data: Dict[str, Any], csvfile) -> None:
        """写入疾病关联靶点CSV数据"""
        import csv
        
        try:
            # 获取所有唯一的datasource.id
            rows = []
            if isinstance(data, dict) and 'data' in data and isinstance(data['data'], dict):
                disease = data['data'].get('disease', {})
                if isinstance(disease, dict):
                    associated_targets = disease.get('associatedTargets', {})
                    if isinstance(associated_targets, dict):
                        rows = associated_targets.get('rows', [])
            
            # 确保rows是列表
            if not isinstance(rows, list):
                rows = []
                
            all_datasources = set()
            for row in rows:
                if isinstance(row, dict):
                    scores = row.get('datasourceScores', [])
                    if isinstance(scores, list):
                        for score in scores:
                            if isinstance(score, dict):
                                all_datasources.add(score.get('componentId', ''))
            
            # 排序数据源ID以确保一致的列顺序
            datasource_columns = sorted(all_datasources)
            
            # 构建表头
            headers = ["target.id", "target.approvedSymbol", "target.approvedName", "score"] + datasource_columns
            writer = csv.writer(csvfile)
            writer.writerow(headers)
            
            # 写入数据
            for row in rows:
                if not isinstance(row, dict):
                    continue
                    
                target = row.get('target', {})
                if not isinstance(target, dict):
                    target = {}
                    
                scores = row.get('datasourceScores', [])
                if not isinstance(scores, list):
                    scores = []
                    
                # 创建数据源分数映射
                score_map = {}
                for score in scores:
                    if isinstance(score, dict):
                        score_map[score.get('componentId', '')] = str(score.get('score', ''))
                
                # 构建行数据
                row_data = [
                    target.get('id', ''),
                    target.get('approvedSymbol', ''),
                    target.get('approvedName', ''),
                    str(row.get('score', ''))
                ] + [score_map.get(ds, '') for ds in datasource_columns]
                
                writer.writerow(row_data)
        except Exception as e:
            logger.error(f"写入疾病关联靶点CSV失败: {str(e)}")
            # 写入错误信息作为CSV的第一行
            writer = csv.writer(csvfile)
            writer.writerow([f"错误: {str(e)}"])


def opentargets_disease_drugs(args):
    """OpenTargets疾病药物查询功能"""
    try:
        config_manager = ConfigManager(args.config)
        ot_manager = OpenTargetsManager(config_manager)
        
        # 获取疾病ID
        if args.disease_id:
            disease_id = args.disease_id
        else:
            disease_id = ot_manager.get_disease_efo_id(
                args.disease_name, 
                interactive=not args.non_interactive
            )
        
        # 查询疾病药物信息
        logger.info(f"开始查询疾病 {disease_id} 的药物信息")
        data = ot_manager.query_disease_drugs(disease_id, args.limit)
        
        # 创建保存目录
        result_dir = os.path.join('result', 'opentargets')
        os.makedirs(result_dir, exist_ok=True)
        
        # 生成文件名
        current_date = datetime.now().strftime("%Y%m%d")
        filename = os.path.join(
            result_dir, 
            f"{disease_id}_disease_drugs_size{args.limit}_{current_date}.json"
        )
        
        # 保存结果
        DataFormatter.save_json(data, filename)
        
        # 转换为CSV格式
        if args.format == 'all' or args.format == 'csv':
            OpenTargetsDataFormatter.convert_to_csv(data, filename, 'disease_drugs')
        
        # ZIP归档将由主函数根据需要创建
        
        print(f"查询完成！结果已保存到 {filename}")
    except BioInfoError as e:
        logger.error(f"错误: {e}")
        print(f"错误: {e}")
        sys.exit(1)


def opentargets_target_associations(args):
    """OpenTargets基因关联疾病查询功能"""
    try:
        config_manager = ConfigManager(args.config)
        ot_manager = OpenTargetsManager(config_manager)
        
        # 获取基因ID
        gene_id = ot_manager.get_ensembl_id(args.gene_name)
        
        # 查询基因关联疾病信息
        logger.info(f"开始查询基因 {gene_id} 关联的疾病")
        data = ot_manager.query_target_associations(gene_id, args.limit)
        
        # 创建保存目录
        result_dir = os.path.join('result', 'opentargets')
        os.makedirs(result_dir, exist_ok=True)
        
        # 生成文件名
        current_date = datetime.now().strftime("%Y%m%d")
        filename = os.path.join(
            result_dir, 
            f"{gene_id}_target_associations_size{args.limit}_{current_date}.json"
        )
        
        # 保存结果
        DataFormatter.save_json(data, filename)
        
        # 转换为CSV格式
        if args.format == 'all' or args.format == 'csv':
            OpenTargetsDataFormatter.convert_to_csv(data, filename, 'target_associations')
        
        # ZIP归档将由主函数根据需要创建
        
        print(f"查询完成！结果已保存到 {filename}")
    except BioInfoError as e:
        logger.error(f"错误: {e}")
        print(f"错误: {e}")
        sys.exit(1)


def opentargets_disease_associations(args):
    """OpenTargets疾病关联靶点查询功能"""
    try:
        config_manager = ConfigManager(args.config)
        ot_manager = OpenTargetsManager(config_manager)
        
        # 获取疾病ID
        if args.disease_id:
            disease_id = args.disease_id
        else:
            disease_id = ot_manager.get_disease_efo_id(
                args.disease_name, 
                interactive=not args.non_interactive
            )
        
        # 查询疾病关联靶点信息
        logger.info(f"开始查询疾病 {disease_id} 关联的靶点")
        data = ot_manager.query_disease_associations(disease_id, args.limit)
        
        # 创建保存目录
        result_dir = os.path.join('result', 'opentargets')
        os.makedirs(result_dir, exist_ok=True)
        
        # 生成文件名
        current_date = datetime.now().strftime("%Y%m%d")
        filename = os.path.join(
            result_dir, 
            f"{disease_id}_disease_associations_size{args.limit}_{current_date}.json"
        )
        
        # 保存结果
        DataFormatter.save_json(data, filename)
        
        # 转换为CSV格式
        if args.format == 'all' or args.format == 'csv':
            OpenTargetsDataFormatter.convert_to_csv(data, filename, 'disease_associations')
        
        # ZIP归档将由主函数根据需要创建
        
        print(f"查询完成！结果已保存到 {filename}")
    except BioInfoError as e:
        logger.error(f"错误: {e}")
        print(f"错误: {e}")
        sys.exit(1)